import streamlit as st
import pandas as pd
from config import Config
from cache import MemoryCache
from solve_file.vector_db_utils import vn
from all_data import judge_name_
from db_setting.base import Base, engine
import time

# Initialize cache and database
cache = MemoryCache()
Base.metadata.create_all(engine)
# Connect to MySQL
vn.connect_to_mysql(
    host=Config.MYSQL_HOST,
    dbname=Config.MYSQL_DATABASE,
    user=Config.MYSQL_USER,
    password=Config.MYSQL_PASSWORD,
    port=Config.MYSQL_PORT
)

# Streamlit page configuration
st.set_page_config(page_title="Data Analysis Demo", layout="wide")

# Title
st.title("Data Analysis Demo")

# Create three columns for the layout
col1, col2, col3 = st.columns(3)

# File Upload Section
with col1:
    st.header("Upload File")
    uploaded_file = st.file_uploader("Choose a CSV or Excel file", type=["csv", "xls", "xlsx"])
    
    if uploaded_file is not None:
        with st.spinner("Processing file..."):
            progress_bar = st.progress(0)
            for i in range(100):
                time.sleep(0.02)  # Simulate processing time
                progress_bar.progress(i + 1)
            
            try:
                if uploaded_file.name.endswith(".csv"):
                    df = pd.read_csv(uploaded_file)
                    judge_name_(uploaded_file.name, df)
                elif uploaded_file.name.endswith((".xls", ".xlsx")):
                    df = pd.read_excel(uploaded_file)
                    judge_name_(uploaded_file.name, df)
                else:
                    st.error("Unsupported file format")
                    st.stop()
                
                st.success("File uploaded successfully!")
                st.write("Preview (First 10 Rows):")
                st.dataframe(df.head(10))
            
            except Exception as e:
                st.error(f"Error: {str(e)}")

# RAG Query Section
with col2:
    st.header("RAG Query")
    question_rag = st.text_input("Enter your RAG question:", key="rag_question")
    if st.button("Submit RAG Query"):
        if question_rag:
            with st.spinner("Querying RAG..."):
                progress_bar = st.progress(0)
                for i in range(100):
                    time.sleep(0.01)  # Simulate query time
                    progress_bar.progress(i + 1)
                
                try:
                    response = vn.ask_with_rag(question_rag, top_k=100)
                    st.success("Query completed!")
                    st.write("Answer:")
                    st.write(response)
                except Exception as e:
                    st.error(f"Error: {str(e)}")
        else:
            st.error("Please enter a question.")

# SQL Query Section
with col3:
    st.header("SQL Query")
    question_sql = st.text_input("Enter your SQL question:", key="sql_question")
    if st.button("Submit SQL Query"):
        if question_sql:
            with st.spinner("Generating and running SQL..."):
                progress_bar = st.progress(0)
                for i in range(100):
                    time.sleep(0.015)  # Simulate SQL processing time
                    progress_bar.progress(i + 1)
                
                try:
                    id = cache.generate_id(question=question_sql)
                    # sql = vn.generate_sql(question=question_sql)
                    sql = vn.generate_sql_(question=question_sql)
                    cache.set(id=id, field='question', value=question_sql)
                    cache.set(id=id, field='sql', value=sql)
                    df = vn.run_sql(sql=sql)
                    cache.set(id=id, field='df', value=df)
                    
                    st.success("Query completed!")
                    st.write("Result (First 10 Rows):")
                    st.dataframe(df.head(10))
                except Exception as e:
                    st.error("Please try rephrasing the question.")
        else:
            st.error("Please enter a question.")

# Footer
st.markdown("---")